# 导入异步IO库
import asyncio

# 导入JSON处理库
import json

# 导入UUID库用于生成唯一会话ID
import uuid

# 从starlette导入应用、路由、响应等相关模块
from starlette.applications import Starlette
from starlette.routing import Route
from starlette.responses import JSONResponse, StreamingResponse
from starlette.requests import Request

# 导入uvicorn用于运行ASGI服务器
import uvicorn


# 定义StreamableHTTPMCPServer类，实现MCP服务器
class StreamableHTTPMCPServer:
    """Streamable HTTP MCP 服务器实现"""

    # 构造函数，初始化服务器状态
    def __init__(self):
        # 是否为无状态模式，默认为有状态
        self.stateless_mode = False  # 可以设置为 False 来启用有状态模式
        # 会话信息存储字典
        self.sessions = {}  # 会话列表
        # 消息队列存储字典
        self.message_queues = {}  # 消息队列

    # 生成会话ID的方法
    def _generate_session_id(self) -> str:
        """生成会话 ID（仅在非无状态模式下使用）"""
        # 如果是有状态模式则生成UUID，否则返回None
        return str(uuid.uuid4()) if not self.stateless_mode else None

    # 处理初始化请求
    async def handle_initialize(self, request: Request):
        """处理初始化请求"""
        # 解析请求体中的JSON数据
        data = await request.json()

        # 判断是否为无状态模式
        if self.stateless_mode:
            # 无状态模式：不创建持久会话
            print(" 无状态模式：初始化请求已确认，但不创建持久会话")
            # 返回初始化确认响应
            return JSONResponse(
                {
                    "jsonrpc": "2.0",
                    "id": data.get("id"),
                    "result": {
                        "protocolVersion": "2024-11-05",
                        "capabilities": {"tools": {}, "resources": {}, "prompts": {}},
                        "serverInfo": {
                            "name": "Streamable HTTP MCP Server",
                            "version": "1.0.0",
                            "mode": "stateless",
                        },
                    },
                }
            )
        else:
            # 有状态模式：创建会话
            session_id = self._generate_session_id()
            # 存储会话信息
            self.sessions[session_id] = {
                "protocol_version": data.get("protocolVersion", "2024-11-05"),
                "capabilities": data.get("capabilities", {}),
                "client_info": data.get("clientInfo", {}),
                "created_at": asyncio.get_event_loop().time(),
            }
            # 为该会话创建消息队列
            self.message_queues[session_id] = asyncio.Queue()

            # 返回初始化响应，包含会话ID
            return JSONResponse(
                {
                    "jsonrpc": "2.0",
                    "id": data.get("id"),
                    "result": {
                        "protocolVersion": "2024-11-05",
                        "capabilities": {"tools": {}, "resources": {}, "prompts": {}},
                        "serverInfo": {
                            "name": "Streamable HTTP MCP Server",
                            "version": "1.0.0",
                            "mode": "stateful",
                            "session_id": session_id,
                        },
                    },
                }
            )

    # 处理工具列表请求
    async def handle_tool_list(self, request: Request):
        """处理工具列表请求"""
        # 解析请求体中的JSON数据
        data = await request.json()

        # 如果是有状态模式，校验会话ID
        if not self.stateless_mode:
            session_id = request.headers.get("X-Session-ID")
            if not session_id or session_id not in self.sessions:
                # 会话无效，返回错误
                return JSONResponse({"error": "Invalid session"}, status_code=400)

        # 工具列表对于无状态和有状态服务器都是一样的
        return JSONResponse(
            {
                "jsonrpc": "2.0",
                "id": data.get("id"),
                "result": {
                    "tools": [
                        {
                            "name": "echo",
                            "description": "Echo back the input (stateless)",
                            "inputSchema": {
                                "type": "object",
                                "properties": {"message": {"type": "string"}},
                                "required": ["message"],
                            },
                        },
                        {
                            "name": "streaming_task",
                            "description": "A task that can stream progress (upgradable to SSE)",
                            "inputSchema": {
                                "type": "object",
                                "properties": {
                                    "duration": {"type": "integer", "default": 5},
                                    "stream": {"type": "boolean", "default": False},
                                },
                            },
                        },
                        {
                            "name": "stateful_counter",
                            "description": "A stateful counter (only works in stateful mode)",
                            "inputSchema": {
                                "type": "object",
                                "properties": {
                                    "action": {
                                        "type": "string",
                                        "enum": ["increment", "decrement", "get"],
                                    }
                                },
                            },
                        },
                    ]
                },
            }
        )

    # 处理工具调用请求
    async def handle_tool_call(self, request: Request):
        """处理工具调用请求"""
        # 解析请求体中的JSON数据
        data = await request.json()
        # 获取要调用的工具名称
        tool_name = data["params"]["name"]

        # 如果是有状态模式，校验会话ID
        if not self.stateless_mode:
            session_id = request.headers.get("X-Session-ID")
            if not session_id or session_id not in self.sessions:
                # 会话无效，返回错误
                return JSONResponse({"error": "Invalid session"}, status_code=400)

        # 处理echo工具
        if tool_name == "echo":
            # 简单工具，直接返回结果（无状态）
            return JSONResponse(
                {
                    "jsonrpc": "2.0",
                    "id": data.get("id"),
                    "result": {
                        "content": [
                            {
                                "type": "text",
                                "text": f"Echo: {data['params']['arguments']['message']}",
                            }
                        ]
                    },
                }
            )

        # 处理streaming_task工具
        elif tool_name == "streaming_task":
            # 流式任务，可以选择是否升级为 SSE
            arguments = data["params"]["arguments"]
            # 获取任务持续时间
            duration = arguments.get("duration", 5)
            # 是否使用流式
            use_streaming = arguments.get("stream", False)

            if use_streaming:
                # 升级为 SSE 流式响应
                return await self._handle_streaming_task(data, duration)
            else:
                # 普通响应，等待任务完成
                result = await self._execute_task_sync(duration)
                return JSONResponse(
                    {
                        "jsonrpc": "2.0",
                        "id": data.get("id"),
                        "result": {
                            "content": [
                                {
                                    "type": "text",
                                    "text": f"Task completed after {duration} seconds: {result}",
                                }
                            ]
                        },
                    }
                )

        # 处理stateful_counter工具
        elif tool_name == "stateful_counter":
            # 如果是无状态模式，返回错误
            if self.stateless_mode:
                return JSONResponse(
                    {
                        "jsonrpc": "2.0",
                        "id": data.get("id"),
                        "error": {
                            "code": -32601,
                            "message": "Stateful counter not available in stateless mode",
                        },
                    }
                )

            # 有状态计数器
            session_id = request.headers.get("X-Session-ID")
            # 获取操作类型
            action = data["params"]["arguments"]["action"]

            # 如果会话中没有计数器，初始化为0
            if "counter" not in self.sessions[session_id]:
                self.sessions[session_id]["counter"] = 0

            # 根据操作类型修改计数器
            if action == "increment":
                self.sessions[session_id]["counter"] += 1
            elif action == "decrement":
                self.sessions[session_id]["counter"] -= 1

            # 获取当前计数器值
            counter_value = self.sessions[session_id]["counter"]

            # 返回计数器操作结果
            return JSONResponse(
                {
                    "jsonrpc": "2.0",
                    "id": data.get("id"),
                    "result": {
                        "content": [
                            {
                                "type": "text",
                                "text": f"Counter {action}ed. Current value: {counter_value}",
                            }
                        ]
                    },
                }
            )

        # 未知工具，返回错误
        return JSONResponse({"error": "Unknown tool"}, status_code=400)

    # 同步执行任务（异步等待一段时间）
    async def _execute_task_sync(self, duration: int) -> str:
        """同步执行任务"""
        # 异步等待指定秒数
        await asyncio.sleep(duration)
        # 返回任务结果字符串
        return f"Task result at {asyncio.get_event_loop().time()}"

    # 处理流式任务（SSE）
    async def _handle_streaming_task(self, data, duration):
        """处理流式任务"""

        # 定义异步生成器，逐步推送进度
        async def stream_progress():
            # 发送开始通知
            yield f"data: {json.dumps({
                'jsonrpc': '2.0',
                'method': 'notifications/progress',
                'params': {
                    'type': 'begin',
                    'title': 'Streaming Task',
                    'message': 'Starting streaming task...'
                }
            })}\n\n"

            # 模拟任务执行，发送进度更新
            for i in range(duration):
                await asyncio.sleep(1)
                # 计算进度百分比
                progress = (i + 1) / duration * 100

                # 发送进度更新通知
                yield f"data: {json.dumps({
                    'jsonrpc': '2.0',
                    'method': 'notifications/progress',
                    'params': {
                        'type': 'update',
                        'title': 'Streaming Task',
                        'message': f'Progress: {progress:.1f}%',
                        'percentage': progress
                    }
                })}\n\n"

            # 发送完成通知
            yield f"data: {json.dumps({
                'jsonrpc': '2.0',
                'method': 'notifications/progress',
                'params': {
                    'type': 'end',
                    'title': 'Streaming Task',
                    'message': 'Streaming task completed!'
                }
            })}\n\n"

            # 发送最终结果
            yield f"data: {json.dumps({
                'jsonrpc': '2.0',
                'id': data.get('id'),
                'result': {
                    'content': [
                        {
                            'type': 'text',
                            'text': f'Streaming task completed after {duration} seconds!'
                        }
                    ]
                }
            })}\n\n"

        # 返回StreamingResponse，SSE流式响应
        return StreamingResponse(
            stream_progress(),
            media_type="text/event-stream",
            headers={"Cache-Control": "no-cache", "Connection": "keep-alive"},
        )

    # 处理统一的消息端点
    async def handle_message_endpoint(self, request: Request):
        """处理统一的消息端点"""
        # 如果是GET请求，建立SSE流
        if request.method == "GET":
            # GET 请求：建立 SSE 流（可选）
            return await self._handle_sse_connection(request)
        # 如果是POST请求，处理MCP消息
        elif request.method == "POST":
            # POST 请求：处理 MCP 消息
            return await self._handle_mcp_message(request)
        # 其他方法不允许
        else:
            return JSONResponse({"error": "Method not allowed"}, status_code=405)

    # 处理SSE连接
    async def _handle_sse_connection(self, request: Request):
        """处理 SSE 连接（可选功能）"""
        # 无状态模式不支持SSE
        if self.stateless_mode:
            return JSONResponse(
                {"error": "SSE connections not supported in stateless mode"},
                status_code=400,
            )

        # 获取会话ID
        session_id = request.headers.get("X-Session-ID")
        # 校验会话ID
        if not session_id or session_id not in self.sessions:
            return JSONResponse({"error": "Invalid session"}, status_code=400)

        # 定义SSE流生成器
        async def sse_stream():
            # 获取该会话的消息队列
            queue = self.message_queues[session_id]

            # 发送连接确认
            yield f"data: {json.dumps({
                'jsonrpc': '2.0',
                'method': 'notifications/connection',
                'params': {
                    'status': 'connected',
                    'session_id': session_id,
                    'mode': 'stateful'
                }
            })}\n\n"

            try:
                while True:
                    # 等待消息，超时30秒
                    message = await asyncio.wait_for(queue.get(), timeout=30.0)
                    # 推送消息
                    yield f"data: {json.dumps(message)}\n\n"
            except asyncio.TimeoutError:
                # 超时则发送心跳保持连接
                yield f"data: {json.dumps({
                    'jsonrpc': '2.0',
                    'method': 'notifications/heartbeat',
                    'params': {'timestamp': asyncio.get_event_loop().time()}
                })}\n\n"

        # 返回StreamingResponse，SSE流式响应
        return StreamingResponse(
            sse_stream(),
            media_type="text/event-stream",
            headers={"Cache-Control": "no-cache", "Connection": "keep-alive"},
        )

    # 处理MCP消息
    async def _handle_mcp_message(self, request: Request):
        """处理 MCP 消息"""
        try:
            # 解析请求体中的JSON数据
            data = await request.json()
            # 获取方法名
            method = data.get("method", "")

            # 根据方法名分发处理
            if method == "initialize":
                return await self.handle_initialize(request)
            elif method == "tools/list":
                return await self.handle_tool_list(request)
            elif method == "tools/call":
                return await self.handle_tool_call(request)
            elif method == "shutdown":
                return await self.handle_shutdown(request)
            else:
                # 未知方法，返回错误
                return JSONResponse(
                    {
                        "jsonrpc": "2.0",
                        "id": data.get("id"),
                        "error": {
                            "code": -32601,
                            "message": f"Method not found: {method}",
                        },
                    }
                )
        except Exception as e:
            # 解析异常，返回错误
            return JSONResponse(
                {
                    "jsonrpc": "2.0",
                    "id": request.query_params.get("id"),
                    "error": {"code": -32700, "message": f"Parse error: {str(e)}"},
                },
                status_code=400,
            )

    # 处理关闭请求
    async def handle_shutdown(self, request: Request):
        """处理关闭请求"""
        # 无状态模式无需清理
        if self.stateless_mode:
            return JSONResponse(
                {
                    "jsonrpc": "2.0",
                    "id": 1,
                    "result": {"message": "Stateless server - no cleanup needed"},
                }
            )

        # 解析请求体中的JSON数据
        data = await request.json()
        # 获取会话ID
        session_id = request.headers.get("X-Session-ID")

        # 删除会话信息
        if session_id in self.sessions:
            del self.sessions[session_id]

        # 删除消息队列
        if session_id in self.message_queues:
            del self.message_queues[session_id]

        # 返回清理完成响应
        return JSONResponse(
            {
                "jsonrpc": "2.0",
                "id": data.get("id"),
                "result": {"message": "Session cleaned up"},
            }
        )


# 创建Starlette应用
def create_app():
    """创建 Starlette 应用"""
    # 实例化服务器对象
    server = StreamableHTTPMCPServer()

    # 定义初始化端点处理函数
    async def initialize(request: Request):
        return await server.handle_initialize(request)

    # 定义工具列表端点处理函数
    async def tool_list(request: Request):
        return await server.handle_tool_list(request)

    # 定义工具调用端点处理函数
    async def tool_call(request: Request):
        return await server.handle_tool_call(request)

    # 定义统一消息端点处理函数
    async def message_endpoint(request: Request):
        return await server.handle_message_endpoint(request)

    # 定义关闭端点处理函数
    async def shutdown(request: Request):
        return await server.handle_shutdown(request)

    # 定义主页端点处理函数
    async def homepage(request: Request):
        return JSONResponse(
            {
                "message": "Streamable HTTP MCP Server",
                "endpoints": {
                    "/initialize": "POST - Initialize MCP connection",
                    "/tools/list": "POST - List available tools",
                    "/tools/call": "POST - Call a tool",
                    "/message": "GET/POST - Unified message endpoint (SSE + MCP)",
                    "/shutdown": "POST - Shutdown connection (stateful mode only)",
                },
                "features": {
                    "stateless_mode": server.stateless_mode,
                    "unified_endpoint": True,
                    "upgradable_requests": True,
                    "flexible_session_management": True,
                },
                "note": "This server demonstrates the advantages of Streamable HTTP transport",
            }
        )

    # 返回Starlette应用对象，注册各路由
    return Starlette(
        routes=[
            Route("/", homepage),
            Route("/message", message_endpoint, methods=["GET", "POST"]),
            Route("/initialize", initialize, methods=["POST"]),
            Route("/tools/list", tool_list, methods=["POST"]),
            Route("/tools/call", tool_call, methods=["POST"]),
            Route("/shutdown", shutdown, methods=["POST"]),
        ]
    )


# 主程序入口
if __name__ == "__main__":
    # 创建应用
    app = create_app()
    # 打印服务器启动信息
    print("启动 Streamable HTTP MCP 服务器...")
    print(" 特点：")
    print("   支持无状态和有状态模式")
    print("   统一的 /message 端点")
    print("   可升级的请求（支持 SSE）")
    print("   灵活的会话管理")
    print("   与现有 HTTP 基础设施完全兼容")
    print("   渐进式功能升级")
    print("服务器地址: http://127.0.0.1:8000")
    print(" 查看 API 文档: http://127.0.0.1:8000")

    # 启动uvicorn服务器
    uvicorn.run(app, host="127.0.0.1", port=8000)
